from dezero.core import *
import numpy as np
from dezero import Variable
'''
注意 _dot_var和_dot_func是私有函数，不能靠import*导入
'''

def _dot_var(v: Variable, verbose=False):
    """
    绘制变量var 
    verbose为True则打印var的shape和dtype 
    格式：id [label="x",color=orange, style=filled] 
    或者：id [label="x: shape dtype",color=orange, style=filled]
    """
    dot_var = '{} [label="{}",color=orange,style=filled]\n'
    
    name = '' if v.name is None else v.name
    if verbose and v.data is not None:
        if v.name is not None:
            name += ': '
        name += str(v.shape) + ' ' + str(v.dtype)
    return dot_var.format(id(v), name)
    
def _dot_func(f: Function):
    """
    绘制函数func 
    格式：id [label="x",color=lightblue, style=filled, shape=box] 
    绘制方向：a -> b 
    """
    dot_func = '{} [label="{}",color=lightblue, style=filled, shape=box]\n'
    txt = dot_func.format(id(f), f.__class__.__name__)
    #绘制从输入到函数的箭头，绘制从函数到输出的箭头
    dot_edge = '{} -> {}\n'
    for x in f.inputs:
        txt += dot_edge.format(id(x), id(f))
    for y in f.outputs:
        txt += dot_edge.format(id(f), id(y())) #y是weakref
    return txt

def get_dot_graph(output: Variable, verbose=False)->str:
    """
    输入： 
        output：Variable 
        verbose：设置输出标签是否带shape和dtype信息 
    输出：Dot格式字符串 
    功能：根据式子输出Dot格式字符串 
    """
    txt = ''
    funcs = []
    seen_set = set()
    
    def add_func(f):
        if f not in seen_set:
            funcs.append(f)
            seen_set.add(f)
            
    add_func(output.creator)
    txt += _dot_var(output, verbose)
    
    while funcs:
        func = funcs.pop()
        txt += _dot_func(func)
        
        for x in func.inputs:
            txt += _dot_var(x,verbose)
            
            if x.creator is not None:
                add_func(x.creator)
    
    return 'digraph g {\n' + txt + '}'
        
import os
import subprocess

def plot_dot_graph(output:Variable, verbose=False, to_file='graph.png'):
    """
    输入： 
        output：Variable 
        verbose：设置输出标签是否带shape和dtype信息 
    输出：根据Dot文件转换的图像 
    功能：先根据式子创建dot文件，再使用dot命令生成png 
    """
    dot_graph = get_dot_graph(output,verbose)
    
    #将dot数据保存到文件
    tmp_dir = os.path.join(os.path.expanduser('~'),'.dezero')
    if not os.path.exists(tmp_dir):
        os.mkdir(tmp_dir)
    graph_path = os.path.join(tmp_dir,'tmp_graph.dot')
    
    with open(graph_path,'w') as f:
        f.write(dot_graph)
        
    #调用dot命令 dot sample.dot -T png -o sample.png
    extension = os.path.splitext(to_file) #扩展名png,pdf等 
    ext = extension[1][1:] #注意上一个输出extension=('graph', '.png')
    cmd = 'dot {} -T {} -o {}'.format(graph_path, ext, to_file)
    subprocess.run(cmd, shell=True)


# =============================================================================
# Utility functions for numpy (numpy magic)
# =============================================================================
def sum_to(x, shape):
    """Sum elements along axes to output an array of a given shape.

    Args:
        x (ndarray): Input array.
        shape:

    Returns:
        ndarray: Output array of the shape.
    """
    ndim = len(shape)
    lead = x.ndim - ndim
    lead_axis = tuple(range(lead))

    axis = tuple([i + lead for i, sx in enumerate(shape) if sx == 1])
    y = x.sum(lead_axis + axis, keepdims=True)
    if lead > 0:
        y = y.squeeze(lead_axis)
    return y


def reshape_sum_backward(gy, x_shape, axis, keepdims):
    """Reshape gradient appropriately for dezero.functions.sum's backward.

    Args:
        gy (dezero.Variable): Gradient variable from the output by backprop.
        x_shape (tuple): Shape used at sum function's forward.
        axis (None or int or tuple of ints): Axis used at sum function's
            forward.
        keepdims (bool): Keepdims used at sum function's forward.

    Returns:
        dezero.Variable: Gradient variable which is reshaped appropriately
    """
    ndim = len(x_shape)
    tupled_axis = axis
    if axis is None:
        tupled_axis = None
    elif not isinstance(axis, tuple):
        tupled_axis = (axis,)

    if not (ndim == 0 or tupled_axis is None or keepdims):
        actual_axis = [a if a >= 0 else a + ndim for a in tupled_axis]
        shape = list(gy.shape)
        for a in sorted(actual_axis):
            shape.insert(a, 1)
    else:
        shape = gy.shape

    gy = gy.reshape(shape)  # reshape
    return gy

def logsumexp(x, axis=1):
    """对softmax函数求log"""
    # xp = cuda.get_array_module(x)
    m = x.max(axis=axis, keepdims=True)
    y = x - m
    np.exp(y, out=y)
    s = y.sum(axis=axis, keepdims=True)
    np.log(s, out=s)
    m += s
    return m


#===========================卷积运用

def get_conv_outsize(input_size, kernel_size, stride, pad):
    """卷积层垂直/水平方向输出数据的大小
    参数：
        input_size:     输入数据的大小
        kernel_size:    过滤器的大小
        stride:         步幅大小
        pad:            外围填充的大小
    输出：垂直/水平方向的大小
    """
    # （输入列长+上下填充-核列长）除以步长 = 走几步 + 第一个位置
    return (input_size + pad*2 - kernel_size) // stride + 1

def get_deconv_outsize(size, k, s, p):
    return s * (size - 1) + k - 2 * p

def pair(x):
    if isinstance(x, int):
        return (x,x)
    elif isinstance(x, tuple):
        assert len(x) == 2
        return x
    else:
        raise ValueError


#===============================
# Download
#===============================
import os
import urllib.request
def show_progress(block_num, block_size, total_size):
    bar_template = "\r[{}] {:.2f}%"

    downloaded = block_num * block_size
    p = downloaded / total_size * 100
    i = int(downloaded / total_size * 30)
    if p >= 100.0: p = 100.0
    if i >= 30: i = 30
    bar = "#" * i + "." * (30 - i)
    print(bar_template.format(bar, p), end='')


cache_dir = os.path.join(os.path.expanduser('~'), '.dezero')


def get_file(url, file_name=None):
    """Download a file from the `url` if it is not in the cache.

    The file at the `url` is downloaded to the `~/.dezero`.

    Args:
        url (str): URL of the file.
        file_name (str): Name of the file. It `None` is specified the original
            file name is used.

    Returns:
        str: Absolute path to the saved file.
    """
    if file_name is None:
        file_name = url[url.rfind('/') + 1:]
    file_path = os.path.join(cache_dir, file_name)

    if not os.path.exists(cache_dir):
        os.mkdir(cache_dir)

    if os.path.exists(file_path):
        return file_path

    print("Downloading: " + file_name)
    try:
        urllib.request.urlretrieve(url, file_path, show_progress)
    except (Exception, KeyboardInterrupt) as e:
        if os.path.exists(file_path):
            os.remove(file_path)
        raise
    print(" Done")

    return file_path
